package com.yuxl.admin.manager.auth.socket.handler;

import com.yuxl.admin.common.socket.config.WsProperties;
import com.yuxl.admin.common.socket.service.WebSocketService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
@Component
@RequiredArgsConstructor
public class WebSocketConnectionHandler implements WebSocketHandler {

    private final WebSocketService webSocketService;
    private final WsProperties wsProperties;
    private final WebSocketMessageHandler messageHandler;

    @Override
    @NonNull
    public Mono<Void> handle(@NonNull WebSocketSession session) {
        log.debug("正在处理WebSocket连接");

        List<String> pathSegments = extractPathSegments(session);
        if (pathSegments.size() != 2) {
            log.warn("无效的WebSocket连接尝试：路径参数不正确");
            return session.close();
        }

        String username = pathSegments.get(0);
        String token = pathSegments.get(1);

        session.getAttributes().put("username", username);
        session.getAttributes().put("token", token);

        AtomicLong lastPongTime = new AtomicLong(System.currentTimeMillis());

        return webSocketService.validateConnection(username, token)
                .flatMap(isValid -> {
                    if (isValid) {
                        webSocketService.addSession(username, session);
                        log.debug("用户 {} 的WebSocket连接已建立", username);

                        Flux<WebSocketMessage> pingFlux = Flux.interval(Duration.ofSeconds(wsProperties.getHeartRate()))
                                .map(i -> {
                                    log.debug("发送PING消息给用户 {}", username);
                                    return session.pingMessage(buffer -> buffer.wrap("PING".getBytes()));
                                });

                        Flux<WebSocketMessage> inboundMessages = session.receive()
                                .doOnNext(message -> {
                                    if (message.getType() == WebSocketMessage.Type.PONG) {
                                        log.debug("收到用户 {} 的PONG消息", username);
                                        lastPongTime.set(System.currentTimeMillis());
                                    }
                                });

                        Mono<Void> timeoutChecker = Flux.interval(Duration.ofSeconds(1))
                                .flatMap(i -> {
                                    if (System.currentTimeMillis() - lastPongTime.get() > wsProperties.getTimeout() * 1000L) {
                                        log.warn("用户 {} 超过{}秒未响应，断开连接", username, wsProperties.getTimeout());
                                        return session.close();
                                    }
                                    return Mono.empty();
                                })
                                .then();

                        Mono<Void> handleMessages = messageHandler.handle(inboundMessages, session, username, token);

                        return Mono.zip(
                                session.send(pingFlux),
                                handleMessages,
                                timeoutChecker
                        ).then();
                    } else {
                        log.warn("无效的WebSocket连接尝试：验证失败");
                        return session.close();
                    }
                })
                .onErrorResume(e -> {
                    log.error("处理WebSocket连接时发生错误", e);
                    return session.close();
                })
                .doFinally(signalType -> {
                    webSocketService.removeSession(username, session);
                    log.debug("用户 {} 的WebSocket连接已关闭", username);
                });
    }

    private List<String> extractPathSegments(WebSocketSession session) {
        String path = session.getHandshakeInfo().getUri().getPath();
        String[] segments = path.split("/");
        return Arrays.asList(segments).subList(3, segments.length);
    }
}